package edu.cornell.cs4740.wsd.tools;

import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import mulan.classifier.MultiLabelOutput;
import mulan.classifier.lazy.MLkNN;
import mulan.classifier.meta.RAkEL;
import mulan.classifier.transformation.LabelPowerset;
import mulan.data.MultiLabelInstances;
import weka.classifiers.trees.J48;
import weka.core.Instance;
import weka.core.Instances;
import edu.cornell.cs4740.wsd.parsing.BigramDistanceContainer;
import edu.cornell.cs4740.wsd.parsing.TrainingData;
import edu.cornell.cs4740.wsd.parsing.WordDistanceContainer;
import edu.cornell.cs4740.wsd.predictions.BaselinePredictor;

public class MulanFormat {
  public static final int RAKEL = 0;
  public static final int MLKNN = 1;

  private final String labelStart = "<label name=\"";
  private final String labelEnd = "\"></label>";
  private final String attribute = "@attribute ";
  private final String testFileSuffix = "test";
  private final String mulanFilePrefix = "mulan-files/";

  private HashMap<String, HashMap<String, Integer>> wordIndex = new HashMap<String, HashMap<String, Integer>>();
  private int windowSize;
  private boolean bigrams;
  private boolean stemsOnly;
  private String wordHeldOut = "";
  private Map<String, List<TrainingData>> training;
  private BaselinePredictor baseline;

  public MulanFormat(Map<String, List<TrainingData>> training, int windowSize,
      boolean bigrams, boolean stemsOnly) throws Exception {
    this.windowSize = windowSize;
    this.bigrams = bigrams;
    this.stemsOnly = stemsOnly;
    this.training = training;
    baseline = new BaselinePredictor(training);

    makeIndex();

    Set<String> keys = training.keySet();

    for (String k : keys) {
      makeArff(k);
      makeXml(k);
    }
  }
  
  public MulanFormat(Map<String, List<TrainingData>> training, int windowSize,
      boolean bigrams, boolean stemsOnly, String holdout) throws Exception {
    this.windowSize = windowSize;
    this.bigrams = bigrams;
    this.stemsOnly = stemsOnly;
    this.training = training;
    this.wordHeldOut = holdout;
    baseline = new BaselinePredictor(training);

    makeIndex();

    Set<String> keys = training.keySet();

    for (String k : keys) {
      makeArff(k);
      makeXml(k);
    }
  }

  public HashMap<String, HashMap<String, Integer>> getIndex() {
    return wordIndex;
  }
  
   private void makeIndex() {
    Set<String> keys = training.keySet();

    for (String k : keys) {
      List<TrainingData> list = training.get(k);
      HashMap<String, Integer> count = new HashMap<String, Integer>();
      HashMap<Integer, Integer> numCount = new HashMap<Integer, Integer>();
      // int nextIndex = 0;

      // give an index number to all the surrounding words or bigrams
      if (bigrams) {
        for (TrainingData item : list) {
          // bigrams
          List<BigramDistanceContainer> preceeding = item
              .getPreceedingBigrams();
          int lenPreceeding = preceeding.size();
          List<BigramDistanceContainer> following = item.getSucceedingBigrams();
          int lenFollowing = following.size();
          for (int i = 0; i < windowSize; i++) {
            int preIndex = lenPreceeding - 1 - i;
            if (preIndex >= 0) {
              BigramDistanceContainer p = preceeding.get(preIndex);
              String pBigram;
              if (stemsOnly) {
                pBigram = p.getStemmedBigramString();
              } else {
                pBigram = p.getBigramString();
              }
              if (pBigram.equals(wordHeldOut)) {
                continue;
              }
              if (count.containsKey(pBigram)) {
                int oldCount = count.get(pBigram);
                count.put(pBigram, oldCount + 1);
                Integer num = numCount.get(oldCount + 1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(oldCount + 1, num + 1);
                Integer oldNum = numCount.get(oldCount) - 1;
                if (oldNum == 0) {
                  numCount.remove(oldCount);
                } else {
                  numCount.put(oldCount, oldNum);
                }
              } else {
                count.put(pBigram, 1);
                Integer num = numCount.get(1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(1, num);
              }
            }

            if (i < lenFollowing) {
              BigramDistanceContainer f = following.get(i);
              String fBigram;
              if (stemsOnly) {
                fBigram = f.getStemmedBigramString();
              } else {
                fBigram = f.getBigramString();
              }
              if (fBigram.equals(wordHeldOut)) {
                continue;
              }
              if (count.containsKey(fBigram)) {
                int oldCount = count.get(fBigram);
                count.put(fBigram, oldCount + 1);
                Integer num = numCount.get(oldCount + 1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(oldCount + 1, num + 1);
                Integer oldNum = numCount.get(oldCount) - 1;
                if (oldNum == 0) {
                  numCount.remove(oldCount);
                } else {
                  numCount.put(oldCount, oldNum);
                }
              } else {
                count.put(fBigram, 1);
                Integer num = numCount.get(1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(1, num);
              }
            }
          }
        }
      } else {
        for (TrainingData item : list) {
          // unigrams
          List<WordDistanceContainer> preceeding = item.getPreceedingWords();
          int lenPreceeding = preceeding.size();
          List<WordDistanceContainer> following = item.getSucceedingWords();
          int lenFollowing = following.size();
          for (int i = 0; i < windowSize; i++) {
            int preIndex = lenPreceeding - 1 - i;
            if (preIndex >= 0) {
              WordDistanceContainer p = preceeding.get(preIndex);
              String pWord;
              if (stemsOnly) {
                pWord = p.getStemmedWord();
              } else {
                pWord = p.getWord();
              }
              if (pWord.equals(wordHeldOut)) {
                continue;
              }
              if (count.containsKey(pWord)) {
                int oldCount = count.get(pWord);
                count.put(pWord, oldCount + 1);
                Integer num = numCount.get(oldCount + 1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(oldCount + 1, num + 1);
                Integer oldNum = numCount.get(oldCount) - 1;
                if (oldNum == 0) {
                  numCount.remove(oldCount);
                } else {
                  numCount.put(oldCount, oldNum);
                }
              } else {
                count.put(pWord, 1);
                Integer num = numCount.get(1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(1, num);
              }
            }

            if (i < lenFollowing) {
              WordDistanceContainer f = following.get(i);
              String fWord;
              if (stemsOnly) {
                fWord = f.getStemmedWord();
              } else {
                fWord = f.getWord();
              }
              if (fWord.equals(wordHeldOut)) {
                continue;
              }
              if (count.containsKey(fWord)) {
                int oldCount = count.get(fWord);
                count.put(fWord, oldCount + 1);
                Integer num = numCount.get(oldCount + 1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(oldCount + 1, num + 1);
                Integer oldNum = numCount.get(oldCount) - 1;
                if (oldNum == 0) {
                  numCount.remove(oldCount);
                } else {
                  numCount.put(oldCount, oldNum);
                }
              } else {
                count.put(fWord, 1);
                Integer num = numCount.get(1);
                num = (num == null) ? 1 : num + 1;
                numCount.put(1, num);
              }
            }
          }
        }
      }
      
      // make index with words that occur more than minCutOff times
      HashMap<String, Integer> index = new HashMap<String, Integer>();
      int nextIndex = 0;
      int minCutOff = 2; // number of times a word must appear for it to be included
      Set<String> countKeys = count.keySet();
      for (String ck : countKeys) {
        int freq = count.get(ck);
        if (freq > minCutOff) {
          index.put(ck, nextIndex);
          nextIndex++;
        }
      }

      wordIndex.put(k, index);
    }

  }

  private String getLine(TrainingData item) {
    if (bigrams) {
      return getLineBigrams(item);
    } else {
      return getLineUnigrams(item);
    }
  }

  private String getLineUnigrams(TrainingData item) {
    // unigrams
    String word = item.getWord().getWord();
    HashMap<String, Integer> mapping = wordIndex.get(word);

    List<WordDistanceContainer> preceeding = item.getPreceedingWords();
    int lenPreceeding = preceeding.size();
    List<WordDistanceContainer> following = item.getSucceedingWords();
    int lenFollowing = following.size();

    // indices of surrounding words
    HashMap<Integer, Integer> surroundings = new HashMap<Integer, Integer>();

    for (int i = 0; i < windowSize; i++) {
      int preIndex = lenPreceeding - 1 - i;
      if (preIndex >= 0) {
        WordDistanceContainer p = preceeding.get(preIndex);
        String pWord;
        if (stemsOnly) {
          pWord = p.getStemmedWord();
        } else {
          pWord = p.getWord();
        }
        if (mapping.containsKey(pWord)) {
          Integer old = surroundings.get(mapping.get(pWord));
          if (old == null) {
            surroundings.put(mapping.get(pWord), 1);
          } else {
            surroundings.put(mapping.get(pWord), old + 1);
          }
        }
      }

      if (i < lenFollowing) {
        WordDistanceContainer f = following.get(i);
        String fWord;
        if (stemsOnly) {
          fWord = f.getStemmedWord();
        } else {
          fWord = f.getWord();
        }
        if (mapping.containsKey(fWord)) {
          Integer old = surroundings.get(mapping.get(fWord));
          if (old == null) {
            surroundings.put(mapping.get(fWord), 1);
          } else {
            surroundings.put(mapping.get(fWord), old + 1);
          }
        }
      }
    }

    StringBuilder line = new StringBuilder();

    // add first word
    if (surroundings.containsKey(0)) {
      line.append(surroundings.get(0));
    } else {
      line.append("0");
    }
    // add remaining words
    for (int i = 1; i < mapping.size(); i++) {
      if (surroundings.containsKey(i)) {
        line.append("," + surroundings.get(i));
      } else {
        line.append(",0");
      }
    }

    // add senses
    List<Boolean> senses = item.getSenseUsage();
    for (int j = 0; j < senses.size(); j++) {
      if (senses.get(j)) {
        line.append(",1");
      } else {
        line.append(",0");
      }
    }

    return line.toString();
  }

  private String getLineBigrams(TrainingData item) {
    // bigrams
    String word = item.getWord().getWord();
    HashMap<String, Integer> mapping = wordIndex.get(word);

    List<BigramDistanceContainer> preceeding = item.getPreceedingBigrams();
    int lenPreceeding = preceeding.size();
    List<BigramDistanceContainer> following = item.getSucceedingBigrams();
    int lenFollowing = following.size();

    // indices of surrounding words
    HashMap<Integer, Integer> surroundings = new HashMap<Integer, Integer>();

    for (int i = 0; i < windowSize; i++) {
      int preIndex = lenPreceeding - 1 - i;
      if (preIndex >= 0) {
        BigramDistanceContainer p = preceeding.get(preIndex);
        String pBigram;
        if (stemsOnly) {
          pBigram = p.getStemmedBigramString();
        } else {
          pBigram = p.getBigramString();
        }
        if (mapping.containsKey(pBigram)) {
          Integer old = surroundings.get(mapping.get(pBigram));
          if (old == null) {
            surroundings.put(mapping.get(pBigram), 1);
          } else {
            surroundings.put(mapping.get(pBigram), old + 1);
          }
        }
      }

      if (i < lenFollowing) {
        BigramDistanceContainer f = following.get(i);
        String fBigram;
        if (stemsOnly) {
          fBigram = f.getStemmedBigramString();
        } else {
          fBigram = f.getBigramString();
        }
        if (mapping.containsKey(fBigram)) {
          Integer old = surroundings.get(mapping.get(fBigram));
          if (old == null) {
            surroundings.put(mapping.get(fBigram), 1);
          } else {
            surroundings.put(mapping.get(fBigram), old + 1);
          }
        }
      }
    }

    StringBuilder line = new StringBuilder();

    // add first bigram
    if (surroundings.containsKey(0)) {
      line.append(surroundings.get(0));
    } else {
      line.append("0");
    }
    // add remaining bigrams
    for (int i = 1; i < mapping.size(); i++) {
      if (surroundings.containsKey(i)) {
        line.append("," + surroundings.get(i));
      } else {
        line.append(",0");
      }
    }

    // add senses
    List<Boolean> senses = item.getSenseUsage();
    for (int j = 0; j < senses.size(); j++) {
      if (senses.get(j)) {
        line.append(",1");
      } else {
        line.append(",0");
      }
    }

    return line.toString();
  }

  /** word, list of predicted senses for each example in the order given */
  public Map<String, List<List<Boolean>>> getPredictions(
      Map<String, List<TrainingData>> test, int type) throws Exception {
    if (type == RAKEL) {
      return getRAkELPredictions(test);
    } else if (type == MLKNN) {
      return getMLKNNPredictions(test);
    }
    return null;
  }

  private Map<String, List<List<Boolean>>> getRAkELPredictions(
      Map<String, List<TrainingData>> test) throws Exception {

    Map<String, List<List<Boolean>>> predictions = new HashMap<String, List<List<Boolean>>>();
    makeTestArff(test);
    Set<String> keyset = test.keySet();

    for (String key : keyset) {
      String arffFilename = mulanFilePrefix + key + ".arff";
      String xmlFilename = mulanFilePrefix + key + ".xml";

      MultiLabelInstances dataset = new MultiLabelInstances(arffFilename,
          xmlFilename);

      RAkEL model = new RAkEL(new LabelPowerset(new J48()));

      model.build(dataset);

      String unlabeledFilename = mulanFilePrefix + key + testFileSuffix
          + ".arff";
      FileReader reader = new FileReader(unlabeledFilename);
      Instances unlabeledData = new Instances(reader);

      int numInstances = unlabeledData.numInstances();

      List<List<Boolean>> predictList = new ArrayList<List<Boolean>>();
      for (int instanceIndex = 0; instanceIndex < numInstances; instanceIndex++) {
        List<Boolean> predict = new ArrayList<Boolean>();
        Instance instance = unlabeledData.instance(instanceIndex);
        MultiLabelOutput output = model.makePrediction(instance);

        boolean someTrue = false;
        boolean[] result = output.getBipartition();
        for (int j = 0; j < result.length; j++) {
          predict.add(j, result[j]);
          someTrue = someTrue || result[j];
        }

        if (!someTrue) {
          TrainingData ex = test.get(key).get(instanceIndex);
          predict = baseline.predict(ex);
        }
        
        predictList.add(instanceIndex, predict);
      }

      predictions.put(key, predictList);
    }

    return predictions;
  }

  private Map<String, List<List<Boolean>>> getMLKNNPredictions(
      Map<String, List<TrainingData>> test) throws Exception {

    Map<String, List<List<Boolean>>> predictions = new HashMap<String, List<List<Boolean>>>();
    makeTestArff(test);
    Set<String> keyset = test.keySet();

    for (String key : keyset) {
      String arffFilename = mulanFilePrefix + key + ".arff";
      String xmlFilename = mulanFilePrefix + key + ".xml";

      MultiLabelInstances dataset = new MultiLabelInstances(arffFilename,
          xmlFilename);

      MLkNN model = new MLkNN();

      model.build(dataset);

      String unlabeledFilename = mulanFilePrefix + key + testFileSuffix
          + ".arff";
      FileReader reader = new FileReader(unlabeledFilename);
      Instances unlabeledData = new Instances(reader);

      int numInstances = unlabeledData.numInstances();

      List<List<Boolean>> predictList = new ArrayList<List<Boolean>>();
      for (int instanceIndex = 0; instanceIndex < numInstances; instanceIndex++) {
        List<Boolean> predict = new ArrayList<Boolean>();
        Instance instance = unlabeledData.instance(instanceIndex);
        MultiLabelOutput output = model.makePrediction(instance);

        boolean someTrue = false;
        boolean[] result = output.getBipartition();
        for (int j = 0; j < result.length; j++) {
          predict.add(j, result[j]);
          someTrue = someTrue || result[j];
        }

        if (!someTrue) {
          TrainingData ex = test.get(key).get(instanceIndex);
          predict = baseline.predict(ex);
        }

        predictList.add(instanceIndex, predict);
      }

      predictions.put(key, predictList);
    }

    return predictions;
  }

  private void makeTestArff(Map<String, List<TrainingData>> test)
      throws Exception {
    Set<String> keyset = test.keySet();

    for (String name : keyset) {
      List<TrainingData> list = test.get(name);

      FileWriter fstream = new FileWriter(mulanFilePrefix + name
          + testFileSuffix + ".arff");
      BufferedWriter out = new BufferedWriter(fstream);
      out.write("@relation " + name);
      out.newLine();
      out.newLine();
      for (int i = 0; i < wordIndex.get(name).size(); i++) {
        out.write(attribute + i + "w numeric");
        out.newLine();
      }
      int numSenses = list.get(0).getSenseUsage().size();
      for (int j = 0; j < numSenses; j++) {
        out.write(attribute + j + " {0,1}");
        out.newLine();
      }

      out.newLine();
      out.write("@data");
      out.newLine();
      out.newLine();

      for (int k = 0; k < list.size(); k++) {
        out.write(getLine(list.get(k)));
        out.newLine();
      }
      out.close();
    }
  }

  private void makeArff(String name) throws Exception {
    List<TrainingData> list = training.get(name);

    FileWriter fstream = new FileWriter(mulanFilePrefix + name + ".arff");
    BufferedWriter out = new BufferedWriter(fstream);
    out.write("@relation " + name);
    out.newLine();
    out.newLine();
    for (int i = 0; i < wordIndex.get(name).size(); i++) {
      out.write(attribute + i + "w numeric");
      out.newLine();
    }
    int numSenses = list.get(0).getSenseUsage().size();
    for (int j = 0; j < numSenses; j++) {
      out.write(attribute + j + " {0,1}");
      out.newLine();
    }

    out.newLine();
    out.write("@data");
    out.newLine();
    out.newLine();

    for (int k = 0; k < list.size(); k++) {
      out.write(getLine(list.get(k)));
      out.newLine();
    }
    out.close();
  }

  private void makeXml(String name) throws Exception {
    FileWriter fstream = new FileWriter(mulanFilePrefix + name + ".xml");
    BufferedWriter out = new BufferedWriter(fstream);
    out.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>");
    out.newLine();
    out.write("<labels xmlns=\"http://mulan.sourceforge.net/labels\">");
    out.newLine();

    int numSenses = training.get(name).get(0).getSenseUsage().size();
    for (int j = 0; j < numSenses; j++) {
      out.write(labelStart + j + labelEnd);
      out.newLine();
    }

    out.write("</labels>");
    out.newLine();
    out.close();
  }
}
